import torch
import matplotlib.pyplot as plt
from utils.meterutils import AverageMeter, accuracy
import time


class Loop:
    def __init__(self, model, train_loader, test_loader, loss_fn, optimizer, device):
        self.model = model
        self.train_loader = train_loader
        self.test_loader = test_loader
        self.loss_fn = loss_fn
        self.optimizer = optimizer
        self.device = device
        self.metric = {"Train Loss": [], "Train Acc": [], "Test Loss": [], "Test Acc": []}

    def train(self, epoch):
        self.model.train()
        EpochTC = AverageMeter()
        DataLoaderTC = AverageMeter()
        TrainLoss = AverageMeter()
        TrainAcc = AverageMeter()
        t0 = time.time()
        for batch_idx, (data, target) in enumerate(self.train_loader):
            data, target = data.to(self.device), target.to(self.device)
            t1 = time.time()

            self.optimizer.zero_grad()
            output = self.model(data)
            loss = self.loss_fn(output, target)
            acc = accuracy(output, target)
            loss.backward()
            self.optimizer.step()

            TrainLoss.update(loss, n=data.size(0))
            TrainAcc.update(acc[0], n=data.size(0))

            DataLoaderTC.update(t1 - t0)
            EpochTC.update(time.time() - t0)
            t0 = time.time()
        self.metric["Train Loss"].append(TrainLoss.avg.item())
        self.metric["Train Acc"].append(TrainAcc.avg.item())
        print("Train Epoch:%d\tEpochTime(DataTime):%.1f(%.1f)\tLoss:%.4f\tAcc:%.4f%%" % (
            epoch, EpochTC.sum, DataLoaderTC.sum, TrainLoss.avg.item(), TrainAcc.avg.item()),end="\t")

    def test(self, epoch):
        with torch.no_grad():
            self.model.eval()
            TestLoss = AverageMeter()
            TestAcc = AverageMeter()
            for data, target in self.test_loader:
                data, target = data.to(self.device), target.to(self.device)

                output = self.model(data)
                loss = self.loss_fn(output, target)
                acc = accuracy(output, target)

                TestLoss.update(loss, n=data.size(0))
                TestAcc.update(acc[0], n=data.size(0))

            self.metric["Test Loss"].append(TestLoss.avg.item())
            self.metric["Test Acc"].append(TestAcc.avg.item())
            print("Test Loss:%.4f\tTest Acc:%.4f%%" % ( TestLoss.avg.item(), TestAcc.avg.item())
                  )

    def show(self):
        x = range(1, self.metric["Test Loss"].__len__() + 1)
        fig, axes = plt.subplots(1, 2)
        axes[0].set_title("Loss")
        axes[0].plot(x, self.metric["Train Loss"], 'g',label="Train")
        axes[0].plot(x, self.metric["Test Loss"], 'r',label="Test")

        axes[1].set_title("Acc")
        axes[1].plot(x, self.metric["Train Acc"], 'g')
        axes[1].plot(x, self.metric["Test Acc"], 'r')
        fig.legend()
        fig.savefig("./result.jpg")
